// Copyright 2020 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

package utils

import (
	"net/http"
	"strconv"

	"github.com/gin-gonic/gin"
	"github.com/jinzhu/gorm"
	"github.com/joomcode/errorx"

	"github.com/pingcap-incubator/tidb-dashboard/pkg/model"
	"github.com/pingcap-incubator/tidb-dashboard/pkg/repo"

	"github.com/pingcap-incubator/tidb-dashboard/pkg/tidb"
)

const (
	// The key that attached the TiDB connection in the gin Context.
	tiDBConnectionKey = "tidb"
	// Http Header key to specify tidb instance
	instanceHeaderKey = "Instance"
)

//SpecifyInstance get specified tidb instance form Http Header
func SpecifyInstance() gin.HandlerFunc {
	return func(c *gin.Context) {
		value := c.Request.Header.Get(instanceHeaderKey)
		var instanceID int
		instanceID, err := strconv.Atoi(value)
		if err != nil {
			instanceID = tidb.NonInstanceID
		}
		c.Set(instanceHeaderKey, instanceID)
		c.Next()
	}
}

// MWConnectTiDB creates a middleware that attaches TiDB connection to the context, according to the identity
// information attached in the context. If a connection cannot be established, subsequent handlers will be skipped
// and errors will be generated.
//
// This middleware must be placed after the `MWAuthRequired()` middleware, otherwise it will panic.
func MWConnectTiDB(tidbClient *tidb.Client) gin.HandlerFunc {
	return func(c *gin.Context) {
		sessionUser := c.MustGet(SessionUserKey).(*SessionUser)
		if sessionUser == nil {
			panic("invalid sessionUser")
		}
		var username string
		var password string

		if sessionUser.IsSQLUser {
			if !sessionUser.HasTiDBAuth {
				// Only TiDBAuth is able to access. Raise error in this case.
				// The error is privilege error instead of authorization error so that user will not be redirected.
				MakeInsufficientPrivilegeError(c)
				c.Abort()
				return
			}
			username = sessionUser.Username
			password = sessionUser.Password
		} else {
			username = tidbClient.DefaultUsername
			password = tidbClient.DefaultPassword
		}

		checkAndSetConnection(c, tidbClient, username, password)

		defer func() {
			closeConnection(c)
		}()

		c.Next()
	}
}

// MWConnectTiDBDefault 设置默认连接(默认数据库)
func MWConnectTiDBDefault(tidbClient *tidb.Client) gin.HandlerFunc {
	return func(c *gin.Context) {
		checkAndSetConnection(c, tidbClient, tidbClient.DefaultUsername, tidbClient.DefaultPassword)
		defer func() {
			closeConnection(c)
		}()
		c.Next()
	}
}

// MWConnectTiDBFromFactory 设置指定连接，从默认数据库中查询出配置后建立连接
func MWConnectTiDBFromFactory(tidbClientFactory *tidb.Factory) gin.HandlerFunc {
	return func(c *gin.Context) {
		instanceID := c.MustGet(instanceHeaderKey).(int)

		//未指定实例在这里直接返回错误
		if instanceID == tidb.NonInstanceID {
			c.Status(http.StatusBadRequest)
			_ = c.Error(tidb.ErrTiDBSpecifyFailed.New("Not specify tidb instance."))
			c.Abort()
			return
		}

		//从默认连接中查询实例信息
		tidbInstance, err := getInstanceFromDefaultDB(c, tidbClientFactory, uint(instanceID))
		if err != nil {
			c.Status(http.StatusInternalServerError)
			_ = c.Error(err)
			c.Abort()
			return
		}

		if tidbInstance == nil {
			c.Status(http.StatusBadRequest)
			_ = c.Error(tidb.ErrTiDBPermissionFailed.New("No permission or not found the instance. instanceID=%d", instanceID))
			c.Abort()
			return
		}

		tidbClient, err := tidbClientFactory.GetOrResisterClient(tidbInstance)
		if err != nil {
			c.Status(http.StatusInternalServerError)
			_ = c.Error(err)
			c.Abort()
			return
		}

		checkAndSetConnection(c, tidbClient, tidbClient.DefaultUsername, tidbClient.DefaultPassword)

		defer func() {
			closeConnection(c)
		}()

		c.Next()
	}
}

func getInstanceFromDefaultDB(c *gin.Context, tidbClientFactory *tidb.Factory, instanceID uint) (*model.TidbInstance, error) {
	defaultClient := tidbClientFactory.GetDefaultClient()
	checkAndSetConnection(c, defaultClient, defaultClient.DefaultUsername, defaultClient.DefaultPassword)
	db := GetTiDBConnection(c)

	defer func() {
		closeConnection(c)
	}()

	//SQL用户和分享码默认访问全部的实例
	userContext := c.MustGet(SessionUserKey)

	if userContext == nil {
		return nil, tidb.ErrTiDBPermissionFailed.New("Missing user info")
	}

	sessionUser := userContext.(*SessionUser)
	if sessionUser.IsShared || sessionUser.IsSQLUser {
		return repo.GetInstance(db, instanceID)
	}

	return repo.GetInstanceByUserID(db, sessionUser.Username, instanceID)
}

func checkAndSetConnection(c *gin.Context, tidbClient *tidb.Client, username string, password string) {
	db, err := tidbClient.OpenSQLConn(username, password)

	if err != nil {
		if errorx.IsOfType(err, tidb.ErrTiDBAuthFailed) {
			// If TiDB conn is ok when login but fail this time, it means TiDB credential has been changed since
			// login. In this case, we return unauthorized error, so that the front-end can let user to login again.
			MakeUnauthorizedError(c)
		} else {
			// For other kind of connection errors, for example, PD goes away, return these errors directly.
			// In front-end we will simply display these errors but not ask user to login again.
			c.Status(http.StatusInternalServerError)
			_ = c.Error(err)
		}
		c.Abort()
		return
	}
	c.Set(tiDBConnectionKey, db)
}

func closeConnection(c *gin.Context) {
	// We allow tiDBConnectionKey to be cleared by `TakeTiDBConnection`.
	dbInContext := c.MustGet(tiDBConnectionKey)
	if dbInContext != nil {
		dbInContext2 := dbInContext.(*gorm.DB)
		if dbInContext2 != nil {
			_ = dbInContext2.Close()
		}
	}
}

// TakeTiDBConnection takes out the TiDB connection stored in the gin context by `MWConnectTiDB` middleware.
// Subsequent handlers in this context cannot access the TiDB connection any more.
//
// The TiDB connection will no longer be closed automatically after all handlers are finished. You must manually
// close the taken out connection.
func TakeTiDBConnection(c *gin.Context) *gorm.DB {
	db := GetTiDBConnection(c)
	c.Set(tiDBConnectionKey, nil)
	return db
}

// GetTiDBConnection gets the TiDB connection stored in the gin context by `MWConnectTiDB` middleware.
//
// The connection will be closed automatically after all handlers are finished. Thus you must not use it outside
// the request lifetime. If you want to extend the lifetime, use `TakeTiDBConnection`.
func GetTiDBConnection(c *gin.Context) *gorm.DB {
	db := c.MustGet(tiDBConnectionKey).(*gorm.DB)
	return db
}
